use std::env;
use std::path::{Path};
use std::sync::Arc;

use crossbeam::queue::SegQueue;
use rand::{distributions::Alphanumeric, Rng};
use indicatif::{ProgressBar, ProgressStyle};
use reqwest::Client;
use tokio::{fs, io::AsyncWriteExt, time::{sleep, Duration}};
use std::process::Command;

// 定义代理trait
trait Proxy: Send + Sync {
    fn apply_to_client(&self, client_builder: reqwest::ClientBuilder) -> reqwest::ClientBuilder;
    fn to_url(&self) -> String;
}

// HTTP代理实现
struct HttpProxy {
    host: String,
    port: u16,
}

impl HttpProxy {
    fn new(host: String, port: u16) -> Self {
        Self { host, port }
    }
}

impl Proxy for HttpProxy {
    fn apply_to_client(&self, client_builder: reqwest::ClientBuilder) -> reqwest::ClientBuilder {
        let proxy_url = format!("http://{}:{}", self.host, self.port);
        client_builder.proxy(reqwest::Proxy::http(&proxy_url).expect("无效的HTTP代理URL"))
    }
    
    fn to_url(&self) -> String {
        format!("http://{}:{}", self.host, self.port)
    }
}

// SOCKS5代理实现
struct Socks5Proxy {
    host: String,
    port: u16,
}

impl Socks5Proxy {
    fn new(host: String, port: u16) -> Self {
        Self { host, port }
    }
}

impl Proxy for Socks5Proxy {
    fn apply_to_client(&self, client_builder: reqwest::ClientBuilder) -> reqwest::ClientBuilder {
        let proxy_url = format!("socks5://{}:{}", self.host, self.port);
        client_builder.proxy(reqwest::Proxy::all(&proxy_url).expect("无效的SOCKS5代理URL"))
    }
    
    fn to_url(&self) -> String {
        format!("socks5://{}:{}", self.host, self.port)
    }
}

// 定义下载器trait
#[async_trait::async_trait]
trait Downloader {
    async fn download(&self, url: &str) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>>;
}

// HTTP下载器实现
struct HttpDownloader {
    client: Client,
}

impl HttpDownloader {
    fn new() -> Self {
        Self {
            client: Client::new(),
        }
    }
    
    fn with_client(client: Client) -> Self {
        Self {
            client,
        }
    }
}

#[async_trait::async_trait]
impl Downloader for HttpDownloader {
    async fn download(&self, url: &str) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>> {
        let response = self.client.get(url).send().await?;
        if response.status().is_success() {
            Ok(response.bytes().await?.to_vec())
        } else {
            Err(format!("HTTP error: {}", response.status()).into())
        }
    }
}

// TS文件下载器trait
#[async_trait::async_trait]
trait TsFileDownloader {
    async fn download_ts_file(
        &self,
        url: &str,
        temp_dir: &str,
        index: usize,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
}

// 带重试机制的TS文件下载器
struct RetryTsFileDownloader<T: Downloader> {
    downloader: T,
    max_retries: u32,
}

impl<T: Downloader> RetryTsFileDownloader<T> {
    fn new(downloader: T, max_retries: u32) -> Self {
        Self {
            downloader,
            max_retries,
        }
    }
}

#[async_trait::async_trait]
impl<T: Downloader + Send + Sync> TsFileDownloader for RetryTsFileDownloader<T> {
    async fn download_ts_file(
        &self,
        url: &str,
        temp_dir: &str,
        index: usize,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let mut retries = 0;
        
        loop {
            match self.download_single_ts_file(url, temp_dir, index).await {
                Ok(_) => return Ok(()),
                Err(e) => {
                    retries += 1;
                    if retries >= self.max_retries {
                        return Err(format!("下载文件 {} 失败，已重试 {} 次: {}", url, self.max_retries, e).into());
                    }
                    eprintln!("下载文件 {} 失败，正在重试 ({}/{}): {}", url, retries, self.max_retries, e);
                    sleep(Duration::from_secs(1)).await;
                }
            }
        }
    }
}

impl<T: Downloader + Send + Sync> RetryTsFileDownloader<T> {
    async fn download_single_ts_file(
        &self,
        url: &str,
        temp_dir: &str,
        index: usize,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let data = self.downloader.download(url).await?;
        let filename = format!("{:05}.ts", index);
        let file_path = Path::new(temp_dir).join(filename);
        let mut file = fs::File::create(file_path).await?;
        file.write_all(&data).await?;
        file.flush().await?;
        Ok(())
    }
}

// M3U8解析器trait
#[async_trait::async_trait]
trait M3U8Parser {
    async fn parse(&self, url: &str) -> Result<Vec<String>, Box<dyn std::error::Error + Send + Sync>>;
}

// HTTP M3U8解析器实现
struct HttpM3U8Parser<T: Downloader> {
    downloader: T,
}

impl<T: Downloader> HttpM3U8Parser<T> {
    fn new(downloader: T) -> Self {
        Self { downloader }
    }
}

#[async_trait::async_trait]
impl<T: Downloader + Send + Sync> M3U8Parser for HttpM3U8Parser<T> {
    async fn parse(&self, url: &str) -> Result<Vec<String>, Box<dyn std::error::Error + Send + Sync>> {
        let content = String::from_utf8(self.downloader.download(url).await?)?;
        let base_url = if let Some(last_slash) = url.rfind('/') {
            &url[..last_slash + 1]
        } else {
            url
        };
        
        let mut ts_urls = Vec::new();
        
        for line in content.lines() {
            let line = line.trim();
            if line.starts_with('#') || line.is_empty() {
                continue;
            }
            
            if line.ends_with(".m3u8") {
                let nested_url = if line.starts_with("http") {
                    line.to_string()
                } else {
                    format!("{}{}", base_url, line)
                };
                
                println!("发现嵌套的m3u8: {}", nested_url);
                let nested_ts_urls = Box::pin(self.parse(&nested_url)).await?;
                ts_urls.extend(nested_ts_urls);
            } else {
                let ts_url = if line.starts_with("http") {
                    line.to_string()
                } else {
                    format!("{}{}", base_url, line)
                };
                ts_urls.push(ts_url);
            }
        }
        
        Ok(ts_urls)
    }
}

// 临时目录管理器trait
trait TempDirManager {
    fn create_temp_directory(&self) -> Result<String, std::io::Error>;
    async fn remove_temp_directory(&self, dir_path: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
}

// 标准临时目录管理器实现
struct StandardTempDirManager;

impl TempDirManager for StandardTempDirManager {
    fn create_temp_directory(&self) -> Result<String, std::io::Error> {
        let rand_string: String = rand::thread_rng()
            .sample_iter(&Alphanumeric)
            .take(5)
            .map(char::from)
            .collect();
        
        let dir_name = format!("temp{}", rand_string);
        std::fs::create_dir_all(&dir_name)?;
        Ok(dir_name)
    }
    
    async fn remove_temp_directory(&self, dir_path: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        tokio::fs::remove_dir_all(dir_path).await?;
        Ok(())
    }
}

// 视频合并器trait
#[async_trait::async_trait]
trait VideoMerger {
    async fn merge(&self, temp_dir: &str, output_file: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
}

// FFmpeg视频合并器实现
struct FFmpegVideoMerger;

#[async_trait::async_trait]
impl VideoMerger for FFmpegVideoMerger {
    async fn merge(&self, temp_dir: &str, output_file: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        println!("正在合并TS文件为MP4...");
        
        // 进入临时ts文件目录
        // 执行ls *.ts | sort | sed "s/^/file '/;s/$/'/" > filelist.txt
        let mut entries = tokio::fs::read_dir(temp_dir).await?;
        let mut ts_files = Vec::new();
        
        while let Some(entry) = entries.next_entry().await? {
            let path = entry.path();
            if let Some(extension) = path.extension() {
                if extension == "ts" {
                    if let Some(file_name) = path.file_name() {
                        ts_files.push(file_name.to_string_lossy().to_string());
                    }
                }
            }
        }
        
        ts_files.sort();
        
        // 创建filelist.txt文件
        let filelist_path = std::path::Path::new(temp_dir).join("filelist.txt");
        let mut filelist_content = String::new();
        for file in ts_files {
            filelist_content.push_str(&format!("file '{}'\n", file));
        }
        tokio::fs::write(&filelist_path, filelist_content).await?;
        
        // 执行ffmpeg -f concat -i filelist.txt -c copy ../output_01.mp4
        let output = Command::new("ffmpeg")
            .arg("-f")
            .arg("concat")
            .arg("-safe")
            .arg("0")
            .arg("-i")
            .arg("filelist.txt")
            .arg("-c")
            .arg("copy")
            .arg("-y")
            .arg(format!("../{}", output_file))
            .current_dir(&temp_dir)
            .output()?;
        
        // 清理filelist.txt文件
        let _ = tokio::fs::remove_file(&filelist_path).await;
        
        if !output.status.success() {
            let error_message = String::from_utf8_lossy(&output.stderr);
            eprintln!("ffmpeg执行失败: {}", error_message);
            return Err(format!("ffmpeg执行失败: {}", error_message).into());
        }
        
        let output_path = std::path::Path::new(temp_dir).parent().unwrap().join(output_file);
        if !output_path.exists() {
            return Err("合并完成但输出文件不存在".into());
        }
        
        println!("合并完成: {}", output_file);
        Ok(())
    }
}

// 下载工作器
async fn download_worker<T: TsFileDownloader + Send + Sync>(
    _worker_id: usize,
    downloader: T,
    queue: Arc<SegQueue<(usize, String)>>,
    temp_dir: String,
    progress_bar: ProgressBar,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    while !queue.is_empty() {
        if let Some((index, url)) = queue.pop() {
            match downloader.download_ts_file(&url, &temp_dir, index).await {
                Ok(_) => {
                    progress_bar.inc(1);
                },
                Err(e) => {
                    progress_bar.abandon_with_message(format!("下载失败: {}", e));
                    return Err(e);
                }
            }
        }
    }
    
    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    let args: Vec<String> = env::args().collect();
    
    // 解析代理参数
    let mut proxy: Option<Box<dyn Proxy + Send + Sync>> = None;
    let mut m3u8_urls_start_index = 1;
    
    // 检查是否有代理参数
    if args.len() >= 3 && (args[1] == "--proxy" || args[1] == "-p") {
        let proxy_str = &args[2];
        proxy = parse_proxy(proxy_str)?;
        m3u8_urls_start_index = 3;
    }
    
    // 解析URL和自定义文件名
    let m3u8_urls_with_names = if args.len() <= m3u8_urls_start_index {
        match fs::read_to_string("urls.txt").await {
            Ok(content) => {
                content.lines()
                    .map(|line| line.trim().to_string())
                    .filter(|line| !line.is_empty() && !line.starts_with("#"))
                    .map(|line| {
                        let parts: Vec<&str> = line.splitn(2, ' ').collect();
                        if parts.len() == 2 {
                            (parts[0].to_string(), Some(parts[1].to_string()))
                        } else {
                            (parts[0].to_string(), None)
                        }
                    })
                    .collect::<Vec<(String, Option<String>)>>()
            },
            Err(_) => {
                print_usage(&args[0]);
                std::process::exit(1);
            }
        }
    } else {
        args[m3u8_urls_start_index..].iter()
            .map(|url| (url.clone(), None))
            .collect()
    };

    if m3u8_urls_with_names.is_empty() {
        eprintln!("未提供任何m3u8 URL");
        std::process::exit(1);
    }

    // 创建带代理的HTTP客户端构建器
    let client_builder = Client::builder();
    let client_builder = if let Some(ref proxy) = proxy {
        println!("使用代理: {}", proxy.to_url());
        proxy.apply_to_client(client_builder)
    } else {
        client_builder
    };
    
    let client = client_builder.build().expect("无法创建HTTP客户端");
    
    let http_downloader = HttpDownloader::new();
    let m3u8_parser = HttpM3U8Parser::new(http_downloader);
    let temp_dir_manager = StandardTempDirManager;
    let video_merger = FFmpegVideoMerger;

    // 逐个串行下载每个m3u8 URL
    for (index, (m3u8_url, custom_name)) in m3u8_urls_with_names.iter().enumerate() {
        println!("开始处理第 {}/{} 个URL: {}", index + 1, m3u8_urls_with_names.len(), m3u8_url);

        // 解析m3u8文件获取所有ts文件链接
        let ts_urls = m3u8_parser.parse(m3u8_url).await?;
        let total_files = ts_urls.len();
        println!("找到 {} 个ts文件", total_files);

        // 创建临时目录
        let temp_dir = temp_dir_manager.create_temp_directory()?;
        println!("创建临时目录: {}", temp_dir);

        // 创建进度条
        let progress_bar = ProgressBar::new(total_files as u64);
        progress_bar.set_style(
            ProgressStyle::default_bar()
                .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta}) {msg}")
                .unwrap()
                .progress_chars("#>-")
        );

        // 创建任务队列
        let queue = Arc::new(SegQueue::new());
        for (index, url) in ts_urls.into_iter().enumerate() {
            queue.push((index, url));
        }

        // 创建多个并发工作器
        let mut workers = vec![];
        let worker_count = 30; // 并发数
        
        for i in 0..worker_count {
            let client = client.clone();
            let queue = Arc::clone(&queue);
            let temp_dir = temp_dir.clone();
            let progress_bar = progress_bar.clone();
            
            let retry_downloader = RetryTsFileDownloader::new(
                HttpDownloader::with_client(client),
                5,
            );
            
            workers.push(tokio::spawn(async move {
                download_worker(i, retry_downloader, queue, temp_dir, progress_bar).await
            }));
        }

        // 等待所有工作器完成
        let mut has_error = false;
        for worker in workers {
            if let Ok(Err(e)) = worker.await {
                eprintln!("下载过程中发生错误: {}", e);
                has_error = true;
            }
        }

        if has_error {
            progress_bar.abandon_with_message("下载过程中发生错误");
            std::process::exit(1);
        } else {
            progress_bar.finish_with_message("所有文件下载完成!");
        }

        // 确定输出文件名
        let output_file_name = if let Some(name) = custom_name {
            if name.ends_with(".mp4") {
                name.clone()
            } else {
                format!("{}.mp4", name)
            }
        } else {
            format!("output_{}.mp4", index + 1)
        };

        // 自动调用ffmpeg合并TS文件为MP4
        match video_merger.merge(&temp_dir, &output_file_name).await {
            Ok(_) => {
                println!("已成功将TS文件合并为MP4: {}", output_file_name);
                // 删除临时目录
                if let Err(e) = temp_dir_manager.remove_temp_directory(&temp_dir).await {
                    eprintln!("删除临时目录 {} 时出错: {}", temp_dir, e);
                } else {
                    println!("已删除临时目录: {}", temp_dir);
                }
            }
            Err(e) => {
                eprintln!("合并TS文件时出错: {}", e);
                println!("您可以手动使用以下命令合并:");
                println!("cd {}", temp_dir);
                println!("ffmpeg -i \"concat:$(ls -1v *.ts | tr '\\n' '|')\" -c copy ../{}", output_file_name);
            }
        }

        println!("\n--------------------------------------------------\n");
    }

    println!("所有m3u8 URL处理完成!");
    Ok(())
}

fn parse_proxy(proxy_str: &str) -> Result<Option<Box<dyn Proxy + Send + Sync>>, Box<dyn std::error::Error + Send + Sync>> {
    if proxy_str.starts_with("http://") {
        let without_prefix = &proxy_str[7..];
        let parts: Vec<&str> = without_prefix.split(':').collect();
        if parts.len() != 2 {
            return Err("HTTP代理格式错误，应为 http://host:port".into());
        }
        let host = parts[0].to_string();
        let port: u16 = parts[1].parse().map_err(|_| "端口号无效")?;
        Ok(Some(Box::new(HttpProxy::new(host, port))))
    } else if proxy_str.starts_with("socks5://") {
        let without_prefix = &proxy_str[9..];
        let parts: Vec<&str> = without_prefix.split(':').collect();
        if parts.len() != 2 {
            return Err("SOCKS5代理格式错误，应为 socks5://host:port".into());
        }
        let host = parts[0].to_string();
        let port: u16 = parts[1].parse().map_err(|_| "端口号无效")?;
        Ok(Some(Box::new(Socks5Proxy::new(host, port))))
    } else {
        Err("不支持的代理类型，支持 http:// 和 socks5://".into())
    }
}

fn print_usage(program_name: &str) {
    eprintln!("使用方法: {} [选项] <m3u8_url1> [m3u8_url2] ... [m3u8_urlN]", program_name);
    eprintln!("选项:");
    eprintln!("  -p, --proxy <proxy_url>  设置代理服务器，支持HTTP和SOCKS5代理");
    eprintln!("                           示例: --proxy http://127.0.0.1:8080");
    eprintln!("                                 --proxy socks5://127.0.0.1:1080");
    eprintln!("");
    eprintln!("或者在当前目录创建urls.txt文件，每行一个m3u8 URL");
}
